package com.artisan.spider.util;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.StringField;
import org.apache.lucene.document.TextField;
import org.apache.lucene.index.*;
import org.apache.lucene.search.*;
import org.apache.lucene.search.highlight.Highlighter;
import org.apache.lucene.search.highlight.QueryScorer;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.QueryBuilder;

import java.io.IOException;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 利用 lucene 创建分词索引、根据关键字查询
 */
public class LuceneUtil {

    /**
     * 创建索引的数据主键
     */
    private String primaryKey;
    /**
     * 索引文件保存位置
     */
    private String indexDirPath;
    /**
     * 不需要分词的键
     */
    private Iterable<String> stringFields;
    /**
     * 需要分词的键
     */
    private Iterable<String> textFields;
    /**
     * 分词器
     */
    private Analyzer analyzer;

    /**
     * @param primaryKey   主键
     * @param indexDirPath 索引创建位置
     * @param stringFields 不需要分词的键
     * @param textFields   需要分词的键
     */
    public LuceneUtil(String primaryKey, String indexDirPath, Iterable<String> stringFields, Iterable<String> textFields) {
        this.primaryKey = primaryKey;
        this.indexDirPath = indexDirPath;
        this.stringFields = stringFields;
        this.textFields = textFields;
        this.analyzer = new StandardAnalyzer();
    }

    /**
     * 创建索引
     * 主键值相同的会被覆盖
     *
     * @param map 索引数据
     * @return
     */
    public long index(Map<String, Object> map) {
        long result = 0;
        IndexWriter indexWriter = null;
        try {
            indexWriter = getIndexWriter();
            Document document = new Document();
            for (String field : stringFields) {
                document.add(new StringField(field, map.get(field) + "", Field.Store.YES));
            }
            for (String field : textFields) {
                document.add(new TextField(field, map.get(field) + "", Field.Store.YES));
            }
            result = indexWriter.updateDocument(new Term(primaryKey, map.get(primaryKey) + ""), document);
            indexWriter.commit();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (indexWriter != null && indexWriter.isOpen()) {
                try {
                    indexWriter.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return result;
    }

    /**
     * 批量创建索引
     * 主键值相同的会被覆盖
     *
     * @param list 索引数据集合
     * @return
     */
    public long index(Iterable<Map<String, Object>> list) {
        long result = 0;
        IndexWriter indexWriter = null;
        try {
            indexWriter = getIndexWriter();
            for (Map<String, Object> map : list) {
                Document document = new Document();
                for (String field : stringFields) {
                    document.add(new StringField(field, map.get(field) + "", Field.Store.YES));
                }
                for (String field : textFields) {
                    document.add(new TextField(field, map.get(field) + "", Field.Store.YES));
                }
                indexWriter.updateDocument(new Term(primaryKey, map.get(primaryKey) + ""), document);
            }
            result = indexWriter.commit();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (indexWriter != null && indexWriter.isOpen()) {
                try {
                    indexWriter.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return result;
    }

    /**
     * 根据主键值删除索引
     *
     * @param primaryKeyValue 主键的值
     * @return
     */
    public long delete(String primaryKeyValue) {
        long result = 0;
        IndexWriter indexWriter = null;
        try {
            indexWriter = getIndexWriter();
            result = indexWriter.deleteDocuments(new Term(primaryKey, primaryKeyValue));
            indexWriter.commit();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (indexWriter != null && indexWriter.isOpen()) {
                try {
                    indexWriter.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return result;
    }


    /**
     * 根据主键值 list 批量删除索引
     *
     * @param primaryKeyValues 主键的值 list
     * @return
     */
    public long delete(Iterable<String> primaryKeyValues) {
        long result = 0;
        IndexWriter indexWriter = null;
        try {
            indexWriter = getIndexWriter();
            for (String value : primaryKeyValues) {
                result += indexWriter.deleteDocuments(new Term(primaryKey, value));
            }
            indexWriter.commit();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (indexWriter != null && indexWriter.isOpen()) {
                try {
                    indexWriter.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return result;
    }

    /**
     * 根据关键字进行查询
     *
     * @param keywords 关键字
     * @param page     分页页数
     * @param limit    分页大小
     * @param map      附加查询条件（必要条件，可以为空）
     * @return
     */
    public Map<String, Object> search(String keywords, int page, int limit, Map<String, Object> map) {
        Map<String, Object> result = new HashMap<>();
        IndexReader indexReader = null;
        try {
            indexReader = getIndexReader();
            IndexSearcher searcher = new IndexSearcher(indexReader);
            BooleanQuery.Builder builder = new BooleanQuery.Builder();
            BooleanQuery.Builder keyWordsBuild = new BooleanQuery.Builder();
            BooleanQuery.Builder mapBuild = new BooleanQuery.Builder();
            QueryBuilder queryBuilder = new QueryBuilder(analyzer);
            for (String field : stringFields) {
                if (map != null && map.get(field) != null && !"".equals(map.get(field))) {
                    String[] values = (map.get(field) + "").split(",");
                    List<BytesRef> bytesRefs = new ArrayList<>();
                    for (String s : values) {
                        bytesRefs.add(new BytesRef(s));
                    }
                    mapBuild.add(new TermInSetQuery(field, bytesRefs), BooleanClause.Occur.MUST);
                }
                if (keywords != null && !"".equals(keywords)) {
                    keyWordsBuild.add(new TermQuery(new Term(field, keywords)), BooleanClause.Occur.SHOULD);
                }
            }
            for (String field : textFields) {
                if (map != null && map.get(field) != null && !"".equals(map.get(field))) {
                    mapBuild.add(queryBuilder.createPhraseQuery(field, map.get(field) + ""), BooleanClause.Occur.MUST);
                }
                if (keywords != null && !"".equals(keywords)) {
                    keyWordsBuild.add(queryBuilder.createPhraseQuery(field, keywords), BooleanClause.Occur.SHOULD);
                }
            }
            if (mapBuild.build().clauses().size() > 0) {
                builder.add(mapBuild.build(), BooleanClause.Occur.MUST);
            }
            if (keyWordsBuild.build().clauses().size() > 0) {
                builder.add(keyWordsBuild.build(), BooleanClause.Occur.MUST);
            }
            TopDocs topDocs = searcher.search(builder.build(), page * limit);
            long count = topDocs.totalHits.value;
            List<Map<String, Object>> list = new ArrayList<>();
            Highlighter highlighter = new Highlighter(new QueryScorer(queryBuilder.createPhraseQuery("", keywords)));
            for (int i = (page - 1) * limit; i < topDocs.scoreDocs.length; i++) {
                Map<String, Object> doc = new HashMap<>();
                Document document = indexReader.document(topDocs.scoreDocs[i].doc);
                for (String field : stringFields) {
                    String highlighterHtml = highlighter.getBestFragment(analyzer, "", document.get(field));
                    doc.put(field, highlighterHtml == null ? document.get(field) : highlighterHtml);
                }
                for (String field : textFields) {
                    String highlighterHtml = highlighter.getBestFragment(analyzer, "", document.get(field));
                    doc.put(field, highlighterHtml == null ? document.get(field) : highlighterHtml);
                }
                list.add(doc);
            }
            result.put("page", page);
            result.put("limit", limit);
            result.put("count", count);
            result.put("data", list);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (indexReader != null) {
                try {
                    indexReader.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return result;
    }

    private IndexWriter getIndexWriter() throws Exception {
        Directory indexDirectory = FSDirectory.open(Paths.get(indexDirPath));
        IndexWriterConfig indexWriterConfig = new IndexWriterConfig(new StandardAnalyzer());
        return new IndexWriter(indexDirectory, indexWriterConfig);

    }

    private IndexReader getIndexReader() throws Exception {
        Directory indexDirectory = FSDirectory.open(Paths.get(indexDirPath));
        return DirectoryReader.open(indexDirectory);
    }

    public Analyzer getAnalyzer() {
        return analyzer;
    }

    public void setAnalyzer(Analyzer analyzer) {
        this.analyzer = analyzer;
    }
}
